////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
//                                                                                                                            //
//                                                         WNProject                                                          //
//                                                                                                                            //
//            This file is distributed under the BSD 2-Clause open source license. See License.txt for details.               //
//                                                                                                                            //
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

#pragma once

#ifndef __WN_SCRIPTING_LLVM_HELPERS_H__
#define __WN_SCRIPTING_LLVM_HELPERS_H__

#define NOMINMAX

#include "WNScripting/inc/WNScriptVariable.h"
#include "WNScripting/inc/WNTypeManager.h"
#include "WNScripting/inc/WNScriptingEngineData.h"
#include "WNScripting/inc/WNDataType.h"
#include "WNScripting/inc/WNFunctionDecl.h"
#include "WNScripting/inc/WNFunctionManager.h"
#include "WNCore/inc/WNPlatform.h"

#ifdef _WN_MSVC
    #pragma warning(push)
    #pragma warning(disable: 4800)
    #pragma warning(disable: 4146)
    #pragma warning(disable: 4355)
    #pragma warning(disable: 4244)
    #pragma warning(disable: 4267)

    #if _WN_MSVC < 160000000
        #pragma warning(disable: 4985)
    #endif
#endif

#include <llvm/IRBuilder.h>

#ifdef _WN_MSVC
    #pragma warning(pop)
#endif

#include <llvm/ExecutionEngine/ExecutionEngine.h>
#include <llvm/Module.h>
#include <llvm/DataLayout.h>
#include <llvm/Support/Casting.h>

namespace WNScripting {
    static WN_INT32 WNLLVMGetTypesPerType(WNScriptingEngineData& _data, llvm::Type* _firstType, llvm::Type* _secondType) {
        const llvm::DataLayout* td = _data.mExecutionEngine->getDataLayout();

        return(static_cast<WN_INT32>(ceill(static_cast<WN_FLOAT64>(td->getTypeAllocSize(_firstType)) / static_cast<WN_FLOAT64>(td->getTypeAllocSize(_secondType)))));
    }

    static WN_UINT64 WNLLVMGetTypeSize(WNScriptingEngineData& _data, llvm::Type* _type) {
        const llvm::DataLayout* td = _data.mExecutionEngine->getDataLayout();
        return(td->getTypeAllocSize(_type));
    }

    static llvm::Value* WNLLVMHelpCreateArrayAlloca(WNDataType* _var, WNScriptingEngineData& _data, WNFunctionDecl* _functionDecl, WN_BOOL _heapAlloc, std::list<WNDataType*>::const_iterator _lengthIt, std::list<WNDataType*>::const_iterator _endIt);

    static llvm::Value* WNLLVMHelpCreateAlloca( WNDataType* _var, WNFunctionDecl* _functionDecl, WNIdentifierElement _identifier, WNScriptingEngineData& _data) {
        llvm::IRBuilder<>* builder = reinterpret_cast<llvm::IRBuilder<>*>(_data.mBuilder);

        if (_data.mTypeManager->IsStructType(_var->GetType())){
            return(builder->CreateAlloca(_data.mTypeManager->GetLLVMType(_var->GetType()), 0, _data.mIdentifierMap->GetStringFromIdentifier(_identifier)));
        } else {
            return(builder->CreateAlloca(_data.mTypeManager->GetLLVMType(_var->GetType()), 0, _data.mIdentifierMap->GetStringFromIdentifier(_identifier)));
        }
    }

    static llvm::Value* WNLLVMHelpCreateStore(llvm::IRBuilderBase& _builder, WNScriptVariable& _var, llvm::Value* _value) {
        return(reinterpret_cast<llvm::IRBuilder<>*>(&_builder)->CreateStore(_value, _var.mLlvmAlloca));
    }

    static llvm::Value* LLVMHelpCreateLoad(llvm::IRBuilderBase& _builder, const WNScriptVariable& _var) {
        return(reinterpret_cast<llvm::IRBuilder<>*>(&_builder)->CreateLoad(_var.mLlvmAlloca));
    }

    static WN_BOOL WNLLVMCreateFree(WNScriptingEngineData& _data, WNFunctionDecl* _functionDecl, llvm::Value* _val, WNTypeElement _elem, WN_UINT32 _numLevels) {
        llvm::IRBuilder<>* builder = reinterpret_cast<llvm::IRBuilder<>*>(_data.mBuilder);
        WNFunction* func = _data.mFunctionManager->GetFreeFunction();
        llvm::Type* intType = _data.mTypeManager->GetIntPtrType();
        llvm::Type* intPtrType = llvm::PointerType::getUnqual(intType);
        llvm::Type* type = _data.mTypeManager->GetLLVMType(_elem);
        llvm::Type* subtype = _data.mTypeManager->GetLLVMType(_data.mTypeManager->GetSubType(_elem));
        std::vector<llvm::Value*> params;

        if (_data.mTypeManager->IsStructType(_elem)) {
            params.push_back(builder->CreatePointerCast(_val, intPtrType));
            builder->CreateCall(func->mLlvmFunction, params, "");
            return(WN_TRUE);
        }

        if (!_data.mTypeManager->IsArrayType(_elem)) {
            return(WN_FALSE);
        }

        if (_numLevels == 0) {
            params.push_back(builder->CreatePointerCast(_val, intPtrType));
            builder->CreateCall(func->mLlvmFunction, params, "");

            return(WN_TRUE);
        } else {
            const llvm::DataLayout* td = _data.mExecutionEngine->getDataLayout();
            WN_UINT64 objsize = td->getTypeAllocSize(type);

            llvm::Value* arrayIntPtr = builder->CreatePointerCast(_val, intPtrType);
            llvm::Value* numElems = builder->CreateLoad(arrayIntPtr, WN_FALSE);
            llvm::Value* numElemsAlloc = builder->CreateAlloca(intType, 0, "numElems");
            llvm::Value* arrayIntPtrAlloc = builder->CreateAlloca(intType, 0, "arrayIntPtr");
            llvm::Value* objectsize = llvm::Constant::getIntegerValue(intType, llvm::APInt(64, objsize));
            llvm::Value* objectsoffset = llvm::Constant::getIntegerValue(intType, llvm::APInt(64, objsize * WNLLVMGetTypesPerType(_data, intType, subtype)));
        
            builder->CreateStore(builder->CreateAdd(_val, objectsoffset), arrayIntPtrAlloc);
            builder->CreateStore(numElems, numElemsAlloc, WN_FALSE);
        
            llvm::BasicBlock* arrHeader = llvm::BasicBlock::Create(llvm::getGlobalContext(), "arrayDeletionHeader", _functionDecl->mLLVMFunction);
            llvm::BasicBlock* arrBody = llvm::BasicBlock::Create(llvm::getGlobalContext(), "arrayDeletionBody", _functionDecl->mLLVMFunction);
            llvm::BasicBlock* doneArr = llvm::BasicBlock::Create(llvm::getGlobalContext(), "arrayDeletionDone", _functionDecl->mLLVMFunction);

            builder->CreateBr(arrHeader);
            builder->SetInsertPoint(arrHeader);
            numElems = builder->CreateLoad(numElemsAlloc, WN_FALSE);

            llvm::Value* cmp = builder->CreateICmpUGT(numElems, llvm::Constant::getIntegerValue(intType, llvm::APInt(64, 0)));

            builder->CreateCondBr(cmp, arrBody, doneArr);
            builder->SetInsertPoint(arrBody);

            arrayIntPtr = builder->CreateLoad(arrayIntPtrAlloc, WN_FALSE);

            llvm::Value* curElemLoc = builder->CreateIntToPtr(arrayIntPtr, subtype);
            llvm::Value* curElem = builder->CreateLoad(curElemLoc);
            WN_BOOL retVal = WNLLVMCreateFree(_data, _functionDecl, curElem, _data.mTypeManager->GetSubType(_elem), _numLevels - 1);

            arrayIntPtr = builder->CreateAdd(arrayIntPtr, objectsize);
            builder->CreateStore(arrayIntPtr, arrayIntPtrAlloc);
            builder->CreateBr(arrHeader);
            builder->SetInsertPoint(doneArr);

            return(retVal);
        }
    }

    static llvm::Value* WNLLVMGenerateAlloc(WNScriptingEngineData& _data, llvm::Type* _type, llvm::Value* _numElems, WN_BOOL _heap) {
        llvm::IRBuilder<>* builder = reinterpret_cast<llvm::IRBuilder<>*>(_data.mBuilder);
        llvm::Value* alloc = WN_NULL;

        if (_heap) {
            WNFunction* func = _data.mFunctionManager->GetMallocFunction();
            const llvm::DataLayout* td = _data.mExecutionEngine->getDataLayout();

            llvm::Type* intPtrType = _data.mTypeManager->GetIntPtrType();
            llvm::Constant* llConst = llvm::ConstantInt::get(intPtrType, td->getTypeAllocSize(_type));

            _numElems = builder->CreateIntCast(_numElems, intPtrType, WN_FALSE);

            llvm::Value* val = builder->CreateMul(llConst, _numElems);
            std::vector<llvm::Value*> params;

            params.push_back(val);

            llvm::Value* retVal = builder->CreateCall(func->mLlvmFunction, params);

            retVal = builder->CreatePointerCast(retVal, llvm::PointerType::getUnqual(_type));

            return(retVal);
        } else {
            return (builder->CreateAlloca(_type, _numElems, "Array!"));
        }
    }

    //TODO:
    // Clean this up a bit :|
    static llvm::Value* WNLLVMHelpCreateArrayAlloca(WNDataType* _var, WNScriptingEngineData& _data, WNFunctionDecl* _functionDecl, WN_BOOL _heapAlloc, std::list<WNDataType*>::const_iterator _lengthIt, std::list<WNDataType*>::const_iterator _endIt) {
        llvm::IRBuilder<>* builder = reinterpret_cast<llvm::IRBuilder<>*>(_data.mBuilder);
        llvm::Value* v = WN_NULL;

        if (_lengthIt != _endIt) {
            v = ((*_lengthIt)->GetOptimizedLength());
            _lengthIt++;
        } else {
            return(WN_NULL);
        }

        llvm::Type * intType = _data.mTypeManager->GetIntPtrType();
        WN_INT32 tempVal = WNLLVMGetTypesPerType(_data, _data.mTypeManager->GetIntPtrType(), _data.mTypeManager->GetLLVMType(_var->GetSubType()->GetType()));
    
        if (v == WN_NULL) {
            static size_t constZero = 0;
            v = llvm::Constant::getIntegerValue(intType, llvm::APInt(64, 0));
        }

        llvm::Value* val = llvm::Constant::getIntegerValue(intType, llvm::APInt(64, tempVal));

        val = builder->CreateIntCast(val, intType, WN_FALSE);
        v = builder->CreateIntCast(v, intType, WN_FALSE);

        llvm::Value* val2 = builder->CreateAdd(val, v);
        llvm::Type* targetType =  intType;
        llvm::Value* alloc = WNLLVMGenerateAlloc(_data, _data.mTypeManager->GetLLVMType(_var->GetSubType()->GetType()), val2, _heapAlloc);

        alloc = builder->CreatePointerCast(alloc, llvm::PointerType::getUnqual(targetType));
        builder->CreateStore(v, alloc, WN_FALSE);
        alloc = builder->CreatePtrToInt(alloc, targetType, "PtrToInt!"); 

        if (_data.mTypeManager->IsArrayType(_var->GetSubType()->GetType())) {
            val = builder->CreateMul(val, llvm::Constant::getIntegerValue(intType, llvm::APInt(64, _data.mExecutionEngine->getDataLayout()->getTypeAllocSize(intType))));
            val = builder->CreateIntCast(val, targetType, WN_FALSE);

            llvm::Value* dataWidth = llvm::Constant::getIntegerValue(intType, llvm::APInt(64, _data.mExecutionEngine->getDataLayout()->getTypeAllocSize(_data.mTypeManager->GetLLVMType(_var->GetSubType()->GetType()))));
            llvm::Value* one = llvm::Constant::getIntegerValue(intType, llvm::APInt(64, 1));
            llvm::Value* tempValLoc = builder->CreateAlloca(targetType, 0, "tempWriteVal");

            builder->CreateStore(builder->CreateIntCast(v, targetType, WN_FALSE), tempValLoc, WN_FALSE);

            llvm::Value* startWriteLoc = builder->CreateAlloca(targetType, 0, "startWriteLoc");
            llvm::Value* startWrite = builder->CreateAdd(alloc, val);

            builder->CreateStore(startWrite, startWriteLoc);

            llvm::BasicBlock* arrBlock = llvm::BasicBlock::Create(llvm::getGlobalContext(), "arrayCreationHeader", _functionDecl->mLLVMFunction);
            llvm::BasicBlock* centerBlock = llvm::BasicBlock::Create(llvm::getGlobalContext(), "arrayCreationBody", _functionDecl->mLLVMFunction);
            llvm::BasicBlock* outBlock = llvm::BasicBlock::Create(llvm::getGlobalContext(), "arrayCreationFooter", _functionDecl->mLLVMFunction);
        
            builder->CreateBr(arrBlock);
            builder->SetInsertPoint(arrBlock);

            llvm::Value* cmp = builder->CreateICmpSLE(builder->CreateLoad(tempValLoc, WN_FALSE, "CheckIfArrayFull"), llvm::ConstantInt::get(targetType, 0, WN_FALSE));

            builder->CreateCondBr(cmp, outBlock, centerBlock);
            builder->SetInsertPoint(centerBlock);
        
            llvm::Value* val = WNLLVMHelpCreateArrayAlloca(_var->GetSubType(), _data, _functionDecl, _heapAlloc, _lengthIt, _endIt);

            if (val != WN_NULL) {
                builder->CreateStore(builder->CreatePtrToInt(val, targetType), builder->CreateIntToPtr(builder->CreateLoad(startWriteLoc), targetType->getPointerTo(0)), WN_TRUE);
            }

            builder->CreateStore(builder->CreateAdd(builder->CreateLoad(startWriteLoc), dataWidth), startWriteLoc);
            builder->CreateStore(builder->CreateSub(builder->CreateLoad(tempValLoc), one), tempValLoc);
            builder->CreateBr(arrBlock);
            builder->SetInsertPoint(outBlock);
        }

        return(alloc);
    }

    static llvm::Value* WNLLVMHelpCreateStructAlloc(WNDataType* _var, WNScriptingEngineData& _data, WN_BOOL _heapAlloc) {
        llvm::IRBuilder<>* builder = reinterpret_cast<llvm::IRBuilder<>*>(_data.mBuilder);
        WNTypeElement elem = _var->GetType();
        llvm::Value* numElems = llvm::ConstantInt::get(_data.mTypeManager->GetIntPtrType(), 1);
        llvm::Value* alloc = WNLLVMGenerateAlloc(_data, _data.mTypeManager->GetLLVMType(elem), numElems, _heapAlloc);

        alloc = builder->CreatePointerCast(alloc, llvm::PointerType::getUnqual(_data.mTypeManager->GetLLVMType(elem)));

        return(alloc);
    }

    static llvm::Value* WNLLVMHelpCreateStructAccess(llvm::Value* _val, WNTypeElement _type, WNScriptingEngineData& _data, WNIdentifierElement _structIdent) {
        WN_DEBUG_ASSERT(_data.mTypeManager->IsStructType(_type));

        return(WN_NULL);
    }
}
#endif // __WN_SCRIPTING_LLVM_HELPERS_H__